#!/usr/bin/env python3
"""
FLUX LoRA训练节点 - 基于基类实现
"""

# 导入ComfyUI模块
try:
    from server import PromptServer
except ImportError:
    PromptServer = None

# 导入基类
from .base_lora_trainer_node import BaseLoRATrainerNode

# 导入参数类型和常量
from .lora_trainer_utils.constants import FLUX_TRAINING_PARAMS, AUX_ID,FLUX_TRAINING_CATEGORY

class SimpleFLUXLoRATrainer(BaseLoRATrainerNode):
    """简单FLUX LoRA训练器节点"""
    
    aux_id = AUX_ID
    
    def __init__(self):
        super().__init__("flux_train_network", "FluxNetworkTrainer")
    
    def _send_progress_event(self, progress_data: dict):
        """FLUX训练器自定义进度事件发送
        
        Args:
            progress_data (dict): 进度数据字典
        """
        try:
            PromptServer.instance.send_sync("flux_training_progress", progress_data)
            print(f"[DEBUG] ✅ FLUX进度事件已发送: {progress_data}")
        except Exception as e:
            print(f"[WARNING] 发送FLUX进度更新失败: {e}")
    
    @classmethod
    def INPUT_TYPES(cls):
        """定义输入类型"""
        return {
            "required": {
                "training_params": (FLUX_TRAINING_PARAMS, {
                    "tooltip": "FLUX训练参数配置"
                }),
            },
            "hidden": {
                "node_id": "UNIQUE_ID"
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("训练状态", "输出路径", "训练日志")
    FUNCTION = "train_simple_flux_lora"
    CATEGORY = FLUX_TRAINING_CATEGORY
    OUTPUT_NODE = True

    def train_simple_flux_lora(self, training_params, node_id=None):
        """简单FLUX LoRA训练"""
        return self.train_lora(training_params, node_id)

# 导出节点
NODE_CLASS_MAPPINGS = {
    "SimpleFLUXLoRATrainer": SimpleFLUXLoRATrainer
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SimpleFLUXLoRATrainer": "🟣 FLUX LoRA训练器"
} 